﻿#region License
// Copyright 2010 Microsoft Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#endregion

using System;
using System.Collections;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Security;
using System.Text;
using MediaWiki.Common;
using Microsoft.Win32;
using MediaWiki.MediaWikiConverter.Utilities.Win32;
using System.Globalization;

namespace MediaWiki.MediaWikiConverter.Utilities.Registry
{
  internal class Win32Registry : IRegistry
  {
    // max size for key/values names that can be fetched
    private const int BufferMaxLength = 1024;

    private static IntPtr GetRegistryKeyHandle(RegistryKeyEx key)
    {
      return key.IsRoot ? new IntPtr((int)key._data) : (IntPtr)key._data;
    }

    private static bool IsHandleValid(RegistryKeyEx key)
    {
      return !key._disposed;
    }

    public object GetValue(RegistryKeyEx rkey, string name, object defaultValue, RegistryValueOptions options)
    {
      RegistryValueKind type = 0;
      int size = 0;
      object obj = null;
      IntPtr handle = GetRegistryKeyHandle(rkey);
      int result = Advapi32.RegQueryValueEx(handle, name, IntPtr.Zero, ref type, IntPtr.Zero, ref size);

      if (result == Win32ResultCode.FileNotFound || result == Win32ResultCode.MarkedForDeletion)
        return defaultValue;

      if (result != Win32ResultCode.MoreData && result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);

      switch (type)
      {
        case RegistryValueKind.String:
          {
            byte[] data;
            result = GetBinaryValue(rkey, name, type, out data, size);
            obj = RegistryKeyEx.DecodeString(data);
          }
          break;
        case RegistryValueKind.ExpandString:
          {
            byte[] data;
            result = GetBinaryValue(rkey, name, type, out data, size);
            obj = RegistryKeyEx.DecodeString(data);
            if ((options & RegistryValueOptions.DoNotExpandEnvironmentNames) == 0)
              obj = Environment.ExpandEnvironmentVariables((string)obj);
          }
          break;
        case RegistryValueKind.Binary:
          {
            byte[] data;
            result = GetBinaryValue(rkey, name, type, out data, size);
            obj = data;
          }
          break;
        case RegistryValueKind.DWord:
          {
            int data = 0;
            result = Advapi32.RegQueryValueEx(handle, name, IntPtr.Zero, ref type, ref data, ref size);
            obj = data;
          }
          break;
        case RegistryValueKind.QWord:
          {
            long data = 0;
            result = Advapi32.RegQueryValueEx(handle, name, IntPtr.Zero, ref type, ref data, ref size);
            obj = data;
          }
          break;
        case RegistryValueKind.MultiString:
          {
            byte[] data;
            result = GetBinaryValue(rkey, name, type, out data, size);

            if (result == Win32ResultCode.Success)
              obj = RegistryKeyEx.DecodeString(data).Split('\0');
          }
          break;
        case RegistryValueKind.Unknown:
          throw new ArgumentOutOfRangeException();
      }

      // check result codes again:
      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);

      return obj;
    }

    private RegistryValueKind CalculateValueKind(object value)
    {
      if (value is int)
      {
        return RegistryValueKind.DWord;
      }
      if (!(value is Array))
      {
        return RegistryValueKind.String;
      }
      if (value is byte[])
      {
        return RegistryValueKind.Binary;
      }
      if (!(value is string[]))
      {
        throw new ArgumentException("{0}.SetValue does not support arrays of type '{1}'. Only Byte[] and String[] are supported.".FormatWith(typeof(RegistryKeyEx).Name, value.GetType().Name));
      }
      return RegistryValueKind.MultiString;
    }

    public void SetValue(RegistryKeyEx rkey, string name, object value)
    {
      SetValue(rkey, name, value, RegistryValueKind.Unknown);
    }

    public void SetValue(RegistryKeyEx rkey, string name, object value, RegistryValueKind valueKind)
    {
      int result;
      IntPtr handle = GetRegistryKeyHandle(rkey);

      if (valueKind == RegistryValueKind.Unknown)
        valueKind = CalculateValueKind(value);

      try
      {
        switch (valueKind)
        {
          case RegistryValueKind.String:
          case RegistryValueKind.ExpandString:
            {
              string rawValue = value.ToStringSafe().NullTerminate();
              result = Advapi32.RegSetValueEx(handle, name, IntPtr.Zero, valueKind, rawValue, rawValue.Length * Marshal.SystemDefaultCharSize);
            }
            break;
          case RegistryValueKind.DWord:
            {
              int rawValue = Convert.ToInt32(value, CultureInfo.InvariantCulture);
              result = Advapi32.RegSetValueEx(handle, name, IntPtr.Zero, RegistryValueKind.DWord, ref rawValue, sizeof(int));
            }
            break;
          case RegistryValueKind.QWord:
            {
              long rawValue = Convert.ToInt64(value, CultureInfo.InvariantCulture);
              result = Advapi32.RegSetValueEx(handle, name, IntPtr.Zero, RegistryValueKind.QWord, ref rawValue, sizeof(long));
            }
            break;
          case RegistryValueKind.MultiString:
            {
              string[] vals = (string[])value;
              StringBuilder fullStringValue = new StringBuilder();
              foreach (string v in vals)
              {
                fullStringValue.Append(v);
                fullStringValue.Append('\0');
              }
              fullStringValue.Append('\0');

              byte[] rawValue = Encoding.Unicode.GetBytes(fullStringValue.ToString());
              result = Advapi32.RegSetValueEx(handle, name, IntPtr.Zero, RegistryValueKind.MultiString, rawValue, rawValue.Length);
            }
            break;
          case RegistryValueKind.Binary:
            {
              byte[] rawValue = (byte[])value;
              result = Advapi32.RegSetValueEx(handle, name, IntPtr.Zero, RegistryValueKind.Binary, rawValue, rawValue.Length);
            }
            break;
          default:
            throw new ArgumentOutOfRangeException("valueKind");
        }
      }
      catch (OverflowException)
      {
        throw new ArgumentException("The type of the value object did not match the specified RegistryValueKind or the object could not be properly converted.");
      }
      catch (InvalidOperationException)
      {
        throw new ArgumentException("The type of the value object did not match the specified RegistryValueKind or the object could not be properly converted.");
      }
      catch (FormatException)
      {
        throw new ArgumentException("The type of the value object did not match the specified RegistryValueKind or the object could not be properly converted.");
      }
      catch (InvalidCastException)
      {
        throw new ArgumentException("The type of the value object did not match the specified RegistryValueKind or the object could not be properly converted.");
      }

      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);
    }

    private int GetBinaryValue(RegistryKeyEx rkey, string name, RegistryValueKind type, out byte[] data, int size)
    {
      byte[] internalData = new byte[size];
      IntPtr handle = GetRegistryKeyHandle(rkey);
      int result = Advapi32.RegQueryValueEx(handle, name, IntPtr.Zero, ref type, internalData, ref size);
      data = internalData;
      return result;
    }

    public int SubKeyCount(RegistryKeyEx rkey)
    {
      StringBuilder stringBuffer = new StringBuilder(BufferMaxLength);
      IntPtr handle = GetRegistryKeyHandle(rkey);

      int index = 0;
      while (true)
      {
        int result = Advapi32.RegEnumKey(handle, index, stringBuffer, stringBuffer.Capacity);

        if (result == Win32ResultCode.NoMoreEntries)
          break;

        if (result != Win32ResultCode.Success)
          GenerateExceptionFromWin32ErrorCode(result);

        index++;
      }

      return index;
    }

    public int ValueCount(RegistryKeyEx rkey)
    {
      StringBuilder buffer = new StringBuilder(BufferMaxLength);
      IntPtr handle = GetRegistryKeyHandle(rkey);

      int index = 0;
      while (true)
      {
        RegistryValueKind type = 0;
        int bufferCapacity = buffer.Capacity;
        int result = Advapi32.RegEnumValue(handle, index, buffer, ref bufferCapacity, IntPtr.Zero, ref type, IntPtr.Zero, IntPtr.Zero);

        if (result == Win32ResultCode.NoMoreEntries)
          break;

        if (result != Win32ResultCode.Success && result != Win32ResultCode.MoreData)
          GenerateExceptionFromWin32ErrorCode(result);

        index++;
      }

      return index;
    }

    public RegistryKeyEx OpenSubKey(RegistryKeyEx rkey, string keyName, bool writable, bool ignoreWow64)
    {
      RegistryAccess access = RegistryAccess.Read;
      if (writable) access |= RegistryAccess.Write;
      if (ignoreWow64) access |= RegistryAccess.WOW64_64Key;

      IntPtr handle = GetRegistryKeyHandle(rkey);

      IntPtr subKeyHandle;
      int result = Advapi32.RegOpenKeyEx(handle, keyName, IntPtr.Zero, access, out subKeyHandle);

      if (result == Win32ResultCode.FileNotFound || result == Win32ResultCode.MarkedForDeletion)
        return null;

      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);

      return new RegistryKeyEx(subKeyHandle, CombineName(rkey, keyName), writable);
    }

    public void Flush(RegistryKeyEx rkey)
    {
      if (!IsHandleValid(rkey))
        return;
      IntPtr handle = GetRegistryKeyHandle(rkey);
      Advapi32.RegFlushKey(handle);
    }

    public void Close(RegistryKeyEx rkey)
    {
      if (!IsHandleValid(rkey))
        return;
      IntPtr handle = GetRegistryKeyHandle(rkey);
      Advapi32.RegCloseKey(handle);
    }

    public RegistryKeyEx CreateSubKey(RegistryKeyEx rkey, string keyName, bool ignoreWow64)
    {
      RegistryAccess access = RegistryAccess.AllAccess;
      if (ignoreWow64) access |= RegistryAccess.WOW64_64Key;

      IntPtr handle = GetRegistryKeyHandle(rkey);
      IntPtr subKeyHandle;
      RegistryResult registryResult;
      int result = Advapi32.RegCreateKeyEx(handle, keyName, IntPtr.Zero, null, RegistryOption.NonVolatile, access, IntPtr.Zero, out subKeyHandle, out registryResult);

      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);

      return new RegistryKeyEx(subKeyHandle, CombineName(rkey, keyName), true);
    }

    public void DeleteKey(RegistryKeyEx rkey, string keyName, bool shouldThrowWhenKeyMissing, bool ignoreWow64)
    {
      RegistryAccess access = 0;
      if (ignoreWow64) access |= RegistryAccess.WOW64_64Key;

      IntPtr handle = GetRegistryKeyHandle(rkey);
      int result = Advapi32.RegDeleteKeyEx(handle, keyName, access, IntPtr.Zero);

      if (result == Win32ResultCode.FileNotFound)
      {
        if (shouldThrowWhenKeyMissing)
          throw new ArgumentException("key " + keyName);

        return;
      }

      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);
    }

    public void DeleteValue(RegistryKeyEx rkey, string value, bool shouldThrowWhenKeyMissing)
    {
      IntPtr handle = GetRegistryKeyHandle(rkey);
      int result = Advapi32.RegDeleteValue(handle, value);

      if (result == Win32ResultCode.MarkedForDeletion)
        return;

      if (result == Win32ResultCode.FileNotFound)
      {
        if (shouldThrowWhenKeyMissing)
          throw new ArgumentException("value " + value);
        return;
      }

      if (result != Win32ResultCode.Success)
        GenerateExceptionFromWin32ErrorCode(result);
    }

    public string[] GetSubKeyNames(RegistryKeyEx rkey)
    {
      IntPtr handle = GetRegistryKeyHandle(rkey);
      StringBuilder buffer = new StringBuilder(BufferMaxLength);
      List<string> keys = new List<string>();

      int index = 0;
      while (true)
      {
        int result = Advapi32.RegEnumKey(handle, index, buffer, buffer.Capacity);

        if (result == Win32ResultCode.NoMoreEntries)
          break;

        if (result != Win32ResultCode.Success)
          GenerateExceptionFromWin32ErrorCode(result);

        keys.Add(buffer.ToString());
        buffer.Length = 0;

        index++;
      }

      return keys.ToArray();
    }

    public string[] GetValueNames(RegistryKeyEx rkey)
    {
      IntPtr handle = GetRegistryKeyHandle(rkey);
      List<string> values = new List<string>();

      int index = 0;
      while (true)
      {
        StringBuilder buffer = new StringBuilder(BufferMaxLength);
        int bufferCapacity = buffer.Capacity;
        RegistryValueKind type = 0;

        int result = Advapi32.RegEnumValue(handle, index, buffer, ref bufferCapacity, IntPtr.Zero, ref type, IntPtr.Zero, IntPtr.Zero);

        if (result == Win32ResultCode.NoMoreEntries)
          break;

        if (result != Win32ResultCode.Success && result != Win32ResultCode.MoreData)
          GenerateExceptionFromWin32ErrorCode(result);

        values.Add(buffer.ToString());

        index++;
      }

      return values.ToArray();
    }

    private void GenerateExceptionFromWin32ErrorCode(int errorCode)
    {
      switch (errorCode)
      {
        case Win32ResultCode.FileNotFound:
        case Win32ResultCode.InvalidParameter:
          throw new ArgumentException();
        case Win32ResultCode.AccessDenied:
          throw new SecurityException();
        case Win32ResultCode.MarkedForDeletion:
          throw RegistryKeyEx.CreateMarkedForDeletionException();
        default:
          // unidentified win32 exception
          throw new Exception();
      }
    }

    /// <summary>
    ///	utility: Combine the sub key name to the current name to produce a 
    ///	fully qualified sub key name.
    /// </summary>
    internal static string CombineName(RegistryKeyEx rkey, string localName)
    {
      return string.Concat(rkey.Name, "\\", localName);
    }
  }
}